package org.apache.solr.client.solrj.response;

import cn.hutool.core.collection.CollUtil;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.LazyDocument;
import org.apache.lucene.document.StoredField;
import org.apache.lucene.index.IndexableField;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrServer;
import org.apache.solr.client.solrj.beans.DocumentObjectBinder;
import org.apache.solr.client.solrj.response.*;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.params.CursorMarkParams;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.common.util.SimpleOrderedMap;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.ResultContext;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.response.transform.DocTransformer;
import org.apache.solr.response.transform.TransformContext;
import org.apache.solr.schema.IndexSchema;
import org.apache.solr.search.DocList;
import org.apache.solr.search.ReturnFields;

import java.io.IOException;
import java.util.*;


/**
 *
 *
 * @since solr 1.3
 */
@SuppressWarnings("unchecked")
public class QueryResponseExt extends SolrResponseBase
{
    // Direct pointers to known types
    private NamedList<Object> _header = null;
    private SolrDocumentList _results = null;
    private NamedList<ArrayList> _sortvalues = null;
    private NamedList<Object> _facetInfo = null;
    private NamedList<Object> _debugInfo = null;
    private NamedList<Object> _highlightingInfo = null;
    private NamedList<NamedList<Object>> _spellInfo = null;
    private NamedList<Object> _statsInfo = null;
    private NamedList<NamedList<Number>> _termsInfo = null;
    private String _cursorMarkNext = null;

    // Grouping response
    private NamedList<Object> _groupedInfo = null;
    private GroupResponse _groupResponse = null;

    private NamedList<Object> _expandedInfo = null;
    private Map<String, SolrDocumentList> _expandedResults = null;

    // Facet stuff
    private Map<String,Integer> _facetQuery = null;
    private List<FacetField> _facetFields = null;
    private List<FacetField> _limitingFacets = null;
    private List<FacetField> _facetDates = null;
    private List<RangeFacet> _facetRanges = null;
    private NamedList<List<PivotField>> _facetPivot = null;
    private List<IntervalFacet> _intervalFacets = null;

    // Highlight Info
    private Map<String,Map<String, Object>> _highlighting = null;

    // SpellCheck Response
    private SpellCheckResponse _spellResponse = null;

    // Terms Response
    private TermsResponse _termsResponse = null;

    // Field stats Response
    private Map<String,FieldStatsInfo> _fieldStatsInfo = null;

    // Debug Info
    private Map<String,Object> _debugMap = null;
    private Map<String,String> _explainMap = null;

    // utility variable used for automatic binding -- it should not be serialized

    public QueryResponseExt(){

    }

    /**
     * Utility constructor to set the solrServer and namedList
     */
    public QueryResponseExt( NamedList<Object> res){
        this.setResponse( res );
    }

    public void setResponse( NamedList<Object> res, SolrQueryRequest req, ReturnFields fields)
    {
        super.setResponse(res);

        // Look for known things
        for( int i=0; i<res.size(); i++ ) {
            String n = res.getName( i );
            if( "responseHeader".equals( n ) ) {
                _header = (NamedList<Object>) res.getVal( i );
            }
            else if( "response".equals( n ) ) {
                _results = extractSolrDocumentList((ResultContext) res.getVal( i ), req, fields);
            }
            else if( "sort_values".equals( n ) ) {
                _sortvalues = (NamedList<ArrayList>) res.getVal( i );
            }
            else if( "facet_counts".equals( n ) ) {
                _facetInfo = (NamedList<Object>) res.getVal( i );
                // extractFacetInfo inspects _results, so defer calling it
                // in case it hasn't been populated yet.
            }
            else if( "debug".equals( n ) ) {
                _debugInfo = (NamedList<Object>) res.getVal( i );
                extractDebugInfo( _debugInfo );
            }
            else if( "grouped".equals( n ) ) {
                _groupedInfo = (NamedList<Object>) res.getVal( i );
                extractGroupedInfo( _groupedInfo );
            }
            else if("expanded".equals(n)) {
                _expandedResults = (Map<String, SolrDocumentList>) res.getVal( i );
            }
            else if( "highlighting".equals( n ) ) {
                _highlightingInfo = (NamedList<Object>) res.getVal( i );
                extractHighlightingInfo( _highlightingInfo );
            }
            else if ( "spellcheck".equals( n ) )  {
                _spellInfo = (NamedList<NamedList<Object>>) res.getVal( i );
                extractSpellCheckInfo( _spellInfo );
            }
            else if ( "stats".equals( n ) )  {
                _statsInfo = (NamedList<Object>) res.getVal( i );
                extractStatsInfo( _statsInfo );
            }
            else if ( "terms".equals( n ) ) {
                _termsInfo = (NamedList<NamedList<Number>>) res.getVal( i );
                extractTermsInfo( _termsInfo );
            }
            else if ( CursorMarkParams.CURSOR_MARK_NEXT.equals( n ) ) {
                _cursorMarkNext = (String) res.getVal( i );
            }
        }
        if(_facetInfo != null) extractFacetInfo( _facetInfo );
    }

    public static SolrDocumentList extractSolrDocumentList(ResultContext res, SolrQueryRequest req, ReturnFields fields) {
        SolrDocumentList result = new SolrDocumentList();
        if(res == null)
            return result;
        DocList ids = res.docs;
        TransformContext context = new TransformContext();
        context.query = res.query;
        context.wantsScores = fields.wantsScore() && ids.hasScores();
        context.req = req;

        result.setStart(ids.offset());
        result.setNumFound(ids.matches());
        result.setMaxScore(context.wantsScores ? new Float(ids.maxScore()) : null);

        DocTransformer transformer = fields.getTransformer();
        context.searcher = req.getSearcher();
        context.iterator = ids.iterator();
        if( transformer != null ) {
            transformer.setContext( context );
        }
        int sz = ids.size();
        Set<String> fnames = fields.getLuceneFieldNames();
        if(CollUtil.isEmpty(fnames))
            fnames = req.getSchema().getFields().keySet();
        try {
            for (int i=0; i<sz; i++) {
                int id = context.iterator.nextDoc();
                Document doc = context.searcher.doc(id, fnames);
                SolrDocument sdoc = toSolrDocument(doc, req.getSchema(), fnames);
                if( transformer != null ) {
                    transformer.transform(sdoc, id);
                }
                result.add(sdoc);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        if( transformer != null ) {
            transformer.setContext( null );
        }
        return result;
    }

    private static SolrDocument toSolrDocument(Document doc, IndexSchema schema, Set<String> fieldNames) throws IOException{
        SolrDocument d = new SolrDocument();
        for (String fname : fieldNames) {
            IndexableField val = doc.getField(fname);
            if ((val instanceof StoredField) || (val instanceof Field) || val instanceof LazyDocument.LazyField) {
                Object o = schema.getField(fname).getType().toObject(val);
                IndexableField[] arr = doc.getFields(fname);
                if(arr != null && arr.length > 1){
                    Object[] valObjs = new Object[arr.length];
                    for (int i = 0; i < valObjs.length; i++) {
                        valObjs[i] = schema.getField(fname).getType().toObject(arr[i]);
                    }
                    d.setField(fname, valObjs);
                }else{
                    d.setField(fname, o);
                }
            }
        }
        return d;
    }
    private void extractSpellCheckInfo(NamedList<NamedList<Object>> spellInfo) {
        _spellResponse = new SpellCheckResponse(spellInfo);
    }

    private void extractTermsInfo(NamedList<NamedList<Number>> termsInfo) {
        _termsResponse = new TermsResponse(termsInfo);
    }

    private void extractStatsInfo(NamedList<Object> info) {
        if( info != null ) {
            _fieldStatsInfo = new HashMap<>();
            NamedList<NamedList<Object>> ff = (NamedList<NamedList<Object>>) info.get( "stats_fields" );
            if( ff != null ) {
                for( Map.Entry<String,NamedList<Object>> entry : ff ) {
                    NamedList<Object> v = entry.getValue();
                    if( v != null ) {
                        _fieldStatsInfo.put( entry.getKey(),
                                new FieldStatsInfo( v, entry.getKey() ) );
                    }
                }
            }
        }
    }

    private void extractDebugInfo( NamedList<Object> debug )
    {
        _debugMap = new LinkedHashMap<>(); // keep the order
        for( Map.Entry<String, Object> info : debug ) {
            _debugMap.put( info.getKey(), info.getValue() );
        }

        // Parse out interesting bits from the debug info
        _explainMap = new HashMap<>();
        NamedList<String> explain = (NamedList<String>)_debugMap.get( "explain" );
        if( explain != null ) {
            for( Map.Entry<String, String> info : explain ) {
                String key = info.getKey();
                _explainMap.put( key, info.getValue() );
            }
        }
    }

    private void extractGroupedInfo( NamedList<Object> info ) {
        if ( info != null ) {
            _groupResponse = new GroupResponse();
            int size = info.size();
            for (int i=0; i < size; i++) {
                String fieldName = info.getName(i);
                Object fieldGroups =  info.getVal(i);
                SimpleOrderedMap<Object> simpleOrderedMap = (SimpleOrderedMap<Object>) fieldGroups;

                Object oMatches = simpleOrderedMap.get("matches");
                Object oNGroups = simpleOrderedMap.get("ngroups");
                Object oGroups = simpleOrderedMap.get("groups");
                Object queryCommand = simpleOrderedMap.get("doclist");
                if (oMatches == null) {
                    continue;
                }

                if (oGroups != null) {
                    Integer iMatches = (Integer) oMatches;
                    ArrayList<Object> groupsArr = (ArrayList<Object>) oGroups;
                    GroupCommand groupedCommand;
                    if (oNGroups != null) {
                        Integer iNGroups = (Integer) oNGroups;
                        groupedCommand = new GroupCommand(fieldName, iMatches, iNGroups);
                    } else {
                        groupedCommand = new GroupCommand(fieldName, iMatches);
                    }

                    for (Object oGrp : groupsArr) {
                        SimpleOrderedMap grpMap = (SimpleOrderedMap) oGrp;
                        Object sGroupValue = grpMap.get( "groupValue");
                        SolrDocumentList doclist = (SolrDocumentList) grpMap.get( "doclist");
                        Group group = new Group(sGroupValue != null ? sGroupValue.toString() : null, doclist) ;
                        groupedCommand.add(group);
                    }

                    _groupResponse.add(groupedCommand);
                } else if (queryCommand != null) {
                    Integer iMatches = (Integer) oMatches;
                    GroupCommand groupCommand;
                    if (oNGroups != null) {
                        Integer iNGroups = (Integer) oNGroups;
                        groupCommand = new GroupCommand(fieldName, iMatches, iNGroups);
                    } else {
                        groupCommand = new GroupCommand(fieldName, iMatches);
                    }
                    SolrDocumentList docList = (SolrDocumentList) queryCommand;
                    groupCommand.add(new Group(fieldName, docList));
                    _groupResponse.add(groupCommand);
                }
            }
        }
    }

    private void extractHighlightingInfo( NamedList<Object> info )
    {
        _highlighting = new HashMap<>();
        for( Map.Entry<String, Object> doc : info ) {
            Map<String, Object> fieldMap = new HashMap<>();
            _highlighting.put( doc.getKey(), fieldMap );

            NamedList<Object> fnl = (NamedList<Object>)doc.getValue();
            for( Map.Entry<String, Object> field : fnl ) {
                fieldMap.put( field.getKey(), field.getValue() );
            }
        }
    }

    private void extractFacetInfo( NamedList<Object> info )
    {
        // Parse the queries
        _facetQuery = new LinkedHashMap<>();
        NamedList<Integer> fq = (NamedList<Integer>) info.get( "facet_queries" );
        if (fq != null) {
            for( Map.Entry<String, Integer> entry : fq ) {
                _facetQuery.put( entry.getKey(), entry.getValue() );
            }
        }

        // Parse the facet info into fields
        // TODO?? The list could be <int> or <long>?  If always <long> then we can switch to <Long>
        NamedList<NamedList<Number>> ff = (NamedList<NamedList<Number>>) info.get( "facet_fields" );
        if( ff != null ) {
            _facetFields = new ArrayList<>( ff.size() );
            _limitingFacets = new ArrayList<>( ff.size() );

            long minsize = _results == null ? Long.MAX_VALUE :_results.getNumFound();
            for( Map.Entry<String,NamedList<Number>> facet : ff ) {
                FacetField f = new FacetField( facet.getKey() );
                for( Map.Entry<String, Number> entry : facet.getValue() ) {
                    f.add( entry.getKey(), entry.getValue().longValue() );
                }

                _facetFields.add( f );
                FacetField nl = f.getLimitingFields( minsize );
                if( nl.getValueCount() > 0 ) {
                    _limitingFacets.add( nl );
                }
            }
        }

        //Parse date facets
        NamedList<NamedList<Object>> df = (NamedList<NamedList<Object>>) info.get("facet_dates");
        if (df != null) {
            // System.out.println(df);
            _facetDates = new ArrayList<>( df.size() );
            for (Map.Entry<String, NamedList<Object>> facet : df) {
                // System.out.println("Key: " + facet.getKey() + " Value: " + facet.getValue());
                NamedList<Object> values = facet.getValue();
                String gap = (String) values.get("gap");
                Date end = (Date) values.get("end");
                FacetField f = new FacetField(facet.getKey(), gap, end);

                for (Map.Entry<String, Object> entry : values)   {
                    try {
                        f.add(entry.getKey(), Long.parseLong(entry.getValue().toString()));
                    } catch (NumberFormatException e) {
                        //Ignore for non-number responses which are already handled above
                    }
                }

                _facetDates.add(f);
            }
        }

        //Parse range facets
        NamedList<NamedList<Object>> rf = (NamedList<NamedList<Object>>) info.get("facet_ranges");
        if (rf != null) {
            _facetRanges = new ArrayList<>( rf.size() );
            for (Map.Entry<String, NamedList<Object>> facet : rf) {
                NamedList<Object> values = facet.getValue();
                Object rawGap = values.get("gap");

                RangeFacet rangeFacet;
                if (rawGap instanceof Number) {
                    Number gap = (Number) rawGap;
                    Number start = (Number) values.get("start");
                    Number end = (Number) values.get("end");

                    Number before = (Number) values.get("before");
                    Number after = (Number) values.get("after");
                    Number between = (Number) values.get("between");

                    rangeFacet = new RangeFacet.Numeric(facet.getKey(), start, end, gap, before, after, between);
                } else {
                    String gap = (String) rawGap;
                    Date start = (Date) values.get("start");
                    Date end = (Date) values.get("end");

                    Number before = (Number) values.get("before");
                    Number after = (Number) values.get("after");
                    Number between = (Number) values.get("between");

                    rangeFacet = new RangeFacet.Date(facet.getKey(), start, end, gap, before, after, between);
                }

                NamedList<Integer> counts = (NamedList<Integer>) values.get("counts");
                for (Map.Entry<String, Integer> entry : counts)   {
                    rangeFacet.addCount(entry.getKey(), entry.getValue());
                }

                _facetRanges.add(rangeFacet);
            }
        }

        //Parse pivot facets
        NamedList pf = (NamedList) info.get("facet_pivot");
        if (pf != null) {
            _facetPivot = new NamedList<>();
            for( int i=0; i<pf.size(); i++ ) {
                _facetPivot.add( pf.getName(i), readPivots( (List<NamedList>)pf.getVal(i) ) );
            }
        }

        //Parse interval facets
        NamedList<NamedList<Object>> intervalsNL = (NamedList<NamedList<Object>>) info.get("facet_intervals");
        if (intervalsNL != null) {
            _intervalFacets = new ArrayList<>(intervalsNL.size());
            for (Map.Entry<String, NamedList<Object>> intervalField : intervalsNL) {
                String field = intervalField.getKey();
                List<IntervalFacet.Count> counts = new ArrayList<IntervalFacet.Count>(intervalField.getValue().size());
                for (Map.Entry<String, Object> interval : intervalField.getValue()) {
                    counts.add(new IntervalFacet.Count(interval.getKey(), (Integer)interval.getValue()));
                }
                _intervalFacets.add(new IntervalFacet(field, counts));
            }
        }
    }

    protected List<PivotField> readPivots( List<NamedList> list )
    {
        ArrayList<PivotField> values = new ArrayList<>( list.size() );
        for( NamedList nl : list ) {
            // NOTE, this is cheating, but we know the order they are written in, so no need to check
            assert "field".equals(nl.getName(0));
            String f = (String)nl.getVal( 0 );
            assert "value".equals(nl.getName(1));
            Object v = nl.getVal( 1 );
            assert "count".equals(nl.getName(2));
            int cnt = ((Integer)nl.getVal( 2 )).intValue();
            List<PivotField> p = null;
            if (4 <= nl.size()) {
                assert "pivot".equals(nl.getName(3));
                Object subPiv = nl.getVal(3);
                assert null != subPiv : "Server sent back 'null' for sub pivots?";
                p = readPivots( (List<NamedList>) subPiv );
            }
            values.add( new PivotField( f, v, cnt, p ) );
        }
        return values;
    }

    //------------------------------------------------------
    //------------------------------------------------------

    /**
     * Remove the field facet info
     */
    public void removeFacets() {
        _facetFields = new ArrayList<>();
    }

    //------------------------------------------------------
    //------------------------------------------------------

    public NamedList<Object> getHeader() {
        return _header;
    }

    public SolrDocumentList getResults() {
        return _results;
    }

    public NamedList<ArrayList> getSortValues(){
        return _sortvalues;
    }

    public Map<String, Object> getDebugMap() {
        return _debugMap;
    }

    public Map<String, String> getExplainMap() {
        return _explainMap;
    }

    public Map<String,Integer> getFacetQuery() {
        return _facetQuery;
    }

    public Map<String, SolrDocumentList> getExpandedResults(){
        return this._expandedResults;
    }

    /**
     * Returns the {@link GroupResponse} containing the group commands.
     * A group command can be the result of one of the following parameters:
     * <ul>
     *   <li>group.field
     *   <li>group.func
     *   <li>group.query
     * </ul>
     *
     * @return the {@link GroupResponse} containing the group commands
     */
    public GroupResponse getGroupResponse() {
        return _groupResponse;
    }

    public Map<String, Map<String, Object>> getHighlighting() {
        return _highlighting;
    }

    public SpellCheckResponse getSpellCheckResponse() {
        return _spellResponse;
    }

    public TermsResponse getTermsResponse() {
        return _termsResponse;
    }

    /**
     * See also: {@link #getLimitingFacets()}
     */
    public List<FacetField> getFacetFields() {
        return _facetFields;
    }

    public List<FacetField> getFacetDates()   {
        return _facetDates;
    }

    public List<RangeFacet> getFacetRanges() {
        return _facetRanges;
    }

    public NamedList<List<PivotField>> getFacetPivot()   {
        return _facetPivot;
    }

    public List<IntervalFacet> getIntervalFacets() {
        return _intervalFacets;
    }

    /** get
     *
     * @param name the name of the
     * @return the FacetField by name or null if it does not exist
     */
    public FacetField getFacetField(String name) {
        if (_facetFields==null) return null;
        for (FacetField f : _facetFields) {
            if (f.getName().equals(name)) return f;
        }
        return null;
    }

    public FacetField getFacetDate(String name)   {
        if (_facetDates == null)
            return null;
        for (FacetField f : _facetDates)
            if (f.getName().equals(name))
                return f;
        return null;
    }

    /**
     * @return a list of FacetFields where the count is less then
     * then #getResults() {@link SolrDocumentList#getNumFound()}
     *
     * If you want all results exactly as returned by solr, use:
     * {@link #getFacetFields()}
     */
    public List<FacetField> getLimitingFacets() {
        return _limitingFacets;
    }

    public <T> List<T> getBeans(Class<T> type){
        return new DocumentObjectBinder().getBeans(type,_results);
    }

    public Map<String, FieldStatsInfo> getFieldStatsInfo() {
        return _fieldStatsInfo;
    }

    public String getNextCursorMark() {
        return _cursorMarkNext;
    }
}



